#include "server.h" // For server_init, server_accept, server_handle_connection, server_cleanup
#include "db.h"     // For init_database, close_database
#include <stdio.h>
#include <stdlib.h>           // For EXIT_FAILURE, EXIT_SUCCESS
#include <signal.h>           // For signal handling
#include <unistd.h>           // For close
#include <errno.h>            // For errno
#include <sys/epoll.h>        // For epoll
#include <fcntl.h>            // For fcntl, F_SETFL, O_NONBLOCK
#include "crypto/sm2_keyex.h" // For sm2_generate_key_pair, sm2_free_key_pair

#define MAX_EPOLL_EVENTS 64 // Max events to handle in one epoll_wait call
#define THREAD_POOL_MIN_THREADS 4
#define THREAD_POOL_MAX_THREADS 8
#define THREAD_POOL_QUEUE_SIZE 100

// Global variable to hold the server socket descriptor for signal handling
static int g_server_fd = -1;
static int g_epoll_fd = -1;
static sqlite3 *g_db = NULL;

// Function to set a file descriptor to non-blocking mode
static int make_socket_non_blocking(int sfd)
{
    int flags, s;
    flags = fcntl(sfd, F_GETFL, 0);
    if (flags == -1)
    {
        perror("fcntl F_GETFL");
        return -1;
    }
    flags |= O_NONBLOCK;
    s = fcntl(sfd, F_SETFL, flags);
    if (s == -1)
    {
        perror("fcntl F_SETFL O_NONBLOCK");
        return -1;
    }
    return 0;
}

void handle_shutdown_signal(int sig)
{
    printf("\nCaught signal %d. Shutting down server gracefully...\n", sig);
    if (g_epoll_fd != -1)
    {
        close(g_epoll_fd);
        g_epoll_fd = -1;
    }
    if (g_server_fd != -1)
    {
        server_cleanup(g_server_fd); // server_cleanup should close the server_fd
        g_server_fd = -1;
    }
    // Client FDs managed by epoll loop will be closed there or already closed.
    // Contexts associated with them should be freed.
    // This is tricky with ongoing threads. Thread pool shutdown should wait for tasks.
    printf("Destroying server thread pool...\n");
    destroy_server_thread_pool(); // Gracefully shutdown thread pool

    if (g_db)
    {
        printf("Closing database connection...\n");
        close_database(g_db);
        g_db = NULL;
    }
    printf("Server shutdown complete.\n");
    exit(EXIT_SUCCESS);
}

/**
 * @brief 文件传输服务器主入口
 * @return int 程序退出状态码
 * @details 初始化数据库连接、线程池和SSL引擎，启动epoll事件循环
 *          处理系统信号(SIGINT/SIGTERM)，实现优雅关闭
 * @warning 需要root权限运行（端口绑定）
 */
int main(int argc, char *argv[])
{
    (void)argc; // Unused for now
    (void)argv; // Unused for now

    // Setup signal handling for graceful shutdown
    signal(SIGINT, handle_shutdown_signal);
    signal(SIGTERM, handle_shutdown_signal);
    signal(SIGPIPE, SIG_IGN); // Ignore SIGPIPE, handle write errors directly

    printf("Initializing database...\n");
    if (init_database(&g_db) != 0)
    {
        fprintf(stderr, "Failed to initialize database.\n");
        return EXIT_FAILURE;
    }
    printf("Database initialized successfully.\n");

    // TODO: Initialize GmSSL/OpenSSL if not handled elsewhere (e.g. in client_init/server_init for crypto parts)

    // Initialize server thread pool (moved from server_init)
    init_server_thread_pool(THREAD_POOL_MIN_THREADS, THREAD_POOL_MAX_THREADS, THREAD_POOL_QUEUE_SIZE);
    if (get_server_thread_pool() == NULL)
    {
        fprintf(stderr, "Failed to initialize server thread pool.\n");
        close_database(g_db);
        return EXIT_FAILURE;
    }

    int port = 8888; // Default port, can be made configurable via argv
    printf("Starting server on port %d...\n", port);
    g_server_fd = server_init(port);
    if (g_server_fd < 0)
    {
        fprintf(stderr, "Failed to initialize server socket.\n");
        destroy_server_thread_pool();
        close_database(g_db);
        return EXIT_FAILURE;
    }

    // server_init likely sets SO_REUSEADDR. Ensure it's non-blocking for epoll.
    if (make_socket_non_blocking(g_server_fd) == -1)
    {
        server_cleanup(g_server_fd);
        destroy_server_thread_pool();
        close_database(g_db);
        return EXIT_FAILURE;
    }

    g_epoll_fd = epoll_create1(0);
    if (g_epoll_fd == -1)
    {
        perror("epoll_create1 failed");
        server_cleanup(g_server_fd);
        destroy_server_thread_pool();
        close_database(g_db);
        return EXIT_FAILURE;
    }

    struct epoll_event event;
    event.data.fd = g_server_fd;
    event.events = EPOLLIN | EPOLLET; // Listen for incoming connections, edge-triggered
    if (epoll_ctl(g_epoll_fd, EPOLL_CTL_ADD, g_server_fd, &event) == -1)
    {
        perror("epoll_ctl EPOLL_CTL_ADD for server_fd failed");
        close(g_epoll_fd);
        server_cleanup(g_server_fd);
        destroy_server_thread_pool();
        close_database(g_db);
        return EXIT_FAILURE;
    }

    struct epoll_event *events = calloc(MAX_EPOLL_EVENTS, sizeof(event));
    if (!events)
    {
        perror("calloc for epoll_events failed");
        close(g_epoll_fd);
        server_cleanup(g_server_fd);
        destroy_server_thread_pool();
        close_database(g_db);
        return EXIT_FAILURE;
    }

    printf("Server listening on port %d with epoll. Press Ctrl+C to shut down.\n", port);

    while (1)
    {
        int n_events = epoll_wait(g_epoll_fd, events, MAX_EPOLL_EVENTS, -1); // Wait indefinitely
        if (n_events == -1)
        {
            if (errno == EINTR)
            { // Interrupted by signal, check if shutdown requested
                continue;
            }
            perror("epoll_wait failed");
            break; // Exit loop on other epoll_wait errors
        }

        for (int i = 0; i < n_events; i++)
        {
            client_context_t *ctx = (client_context_t *)events[i].data.ptr; // For client events
            int current_fd = -1;
            if (events[i].data.fd == g_server_fd)
            { // Special handling for listening socket
                current_fd = g_server_fd;
            }
            else if (ctx)
            { // For client sockets, fd is in context
                current_fd = ctx->client_fd;
            }
            else
            { // Should not happen if ptr is always set for client events
                fprintf(stderr, "epoll event for unknown fd or null context pointer\n");
                // If events[i].data.fd is valid, try to close it directly? Risky.
                // This indicates a logic error in how data.ptr is managed.
                // We might try to get fd directly from events[i].data.fd if ctx is NULL
                // but then we can't do context-specific cleanup.
                // For now, if ctx is NULL for a non-server_fd event, we have an issue.
                // Let's assume if it's not server_fd, data.ptr is always a valid ctx.
                current_fd = events[i].data.fd; // Fallback, but problematic for cleanup
            }

            if ((events[i].events & EPOLLERR) || (events[i].events & EPOLLHUP))
            {
                fprintf(stderr, "epoll error/hup on fd %d.\n", current_fd);
                if (current_fd != g_server_fd && ctx)
                {                                                               // Error on client socket
                    epoll_ctl(g_epoll_fd, EPOLL_CTL_DEL, ctx->client_fd, NULL); // Remove from epoll
                    close(ctx->client_fd);
                    if (ctx->upload_file_ptr)
                        fclose(ctx->upload_file_ptr);
                    if (ctx->download_file_ptr)
                        fclose(ctx->download_file_ptr);
                    sm2_free_key_pair(&ctx->sm2_server_key);
                    free(ctx);
                }
                else if (current_fd != g_server_fd && !ctx)
                {
                    // Error on a client FD but context is null - best effort cleanup
                    epoll_ctl(g_epoll_fd, EPOLL_CTL_DEL, current_fd, NULL);
                    close(current_fd);
                }
                // No specific action for server_fd error here, main loop exit or signal handles it.
                continue;
            }

            if (events[i].data.fd == g_server_fd)
            { // New connection on listening socket
                while (1)
                {
                    int client_fd = server_accept(g_server_fd);
                    if (client_fd == -1)
                    {
                        if ((errno == EAGAIN) || (errno == EWOULDBLOCK))
                            break; // All pending accepted
                        perror("server_main: server_accept failed");
                        break;
                    }
                    // server_accept now sets client_fd to non-blocking.
                    client_context_t *new_ctx = (client_context_t *)calloc(1, sizeof(client_context_t));
                    if (!new_ctx)
                    {
                        perror("calloc for client_context_t failed");
                        close(client_fd);
                        continue;
                    }

                    new_ctx->client_fd = client_fd;
                    new_ctx->epoll_fd = g_epoll_fd;
                    new_ctx->db = g_db;
                    new_ctx->marked_for_close = 0;
                    new_ctx->user_id = -1; // Initialize other fields to default
                    new_ctx->auth_done = 0;
                    new_ctx->key_exchange_done = 0;
                    // Initialize SM2 key pair for this session for the server
                    if (sm2_generate_key_pair(&new_ctx->sm2_server_key) != 0)
                    {
                        fprintf(stderr, "Failed to generate SM2 key pair for new client context (fd:%d)\n", client_fd);
                        free(new_ctx);
                        close(client_fd);
                        continue; // Skip adding this client
                    }

                    struct epoll_event client_event;
                    client_event.data.ptr = new_ctx; // Store pointer to context
                    client_event.events = EPOLLIN | EPOLLET | EPOLLONESHOT;
                    if (epoll_ctl(g_epoll_fd, EPOLL_CTL_ADD, client_fd, &client_event) == -1)
                    {
                        perror("epoll_ctl EPOLL_CTL_ADD for client_fd failed");
                        sm2_free_key_pair(&new_ctx->sm2_server_key);
                        free(new_ctx);
                        close(client_fd);
                    }
                }
            }
            else
            { // Data from an existing client connection
                // ctx should be valid here as populated by EPOLL_CTL_ADD or EPOLL_CTL_MOD
                if (!ctx)
                {
                    fprintf(stderr, "Error: Null context for client event on fd %d\n", events[i].data.fd);
                    // This fd is likely already dead or removed, or a logic error occurred.
                    // Try to remove from epoll and close if fd seems valid from event.
                    epoll_ctl(g_epoll_fd, EPOLL_CTL_DEL, events[i].data.fd, NULL);
                    close(events[i].data.fd);
                    continue;
                }

                if (ctx->marked_for_close)
                { // Marked by process_client_event for cleanup
                    fprintf(stdout, "Cleaning up context for fd %d (marked for close)\n", ctx->client_fd);
                    epoll_ctl(g_epoll_fd, EPOLL_CTL_DEL, ctx->client_fd, NULL);
                    close(ctx->client_fd);
                    if (ctx->upload_file_ptr)
                    {
                        fclose(ctx->upload_file_ptr);
                        ctx->upload_file_ptr = NULL;
                        add_audit_log(ctx->db, ctx->user_id, "UPLOAD_ABORTED", ctx->current_upload_server_path, ERR_CONNECTION_ERROR);
                    }
                    if (ctx->download_file_ptr)
                    {
                        fclose(ctx->download_file_ptr);
                        ctx->download_file_ptr = NULL;
                        add_audit_log(ctx->db, ctx->user_id, "DOWNLOAD_ABORTED", ctx->current_download_server_path, ERR_CONNECTION_ERROR);
                    }
                    sm2_free_key_pair(&ctx->sm2_server_key);
                    free(ctx);
                }
                else
                {
                    // Dispatch to thread pool. process_client_event will re-arm EPOLLONESHOT or mark for close.
                    if (thread_pool_add_task(get_server_thread_pool(), process_client_event, ctx) != 0)
                    {
                        perror("Failed to add client event task to thread pool");
                        // If task submission fails, we need to decide how to handle the client.
                        // For now, mark for close and let epoll loop clean it up next time (or here).
                        ctx->marked_for_close = 1; // Mark it, will be cleaned up in next epoll iteration or here
                        // Potentially clean up immediately to avoid leaving client hanging with no processing
                        fprintf(stderr, "Cleaning up context for fd %d due to thread pool submission failure\n", ctx->client_fd);
                        epoll_ctl(g_epoll_fd, EPOLL_CTL_DEL, ctx->client_fd, NULL);
                        close(ctx->client_fd);
                        if (ctx->upload_file_ptr)
                        {
                            fclose(ctx->upload_file_ptr);
                            ctx->upload_file_ptr = NULL;
                            add_audit_log(ctx->db, ctx->user_id, "UPLOAD_ABORTED", ctx->current_upload_server_path, ERR_CONNECTION_ERROR);
                        }
                        if (ctx->download_file_ptr)
                        {
                            fclose(ctx->download_file_ptr);
                            ctx->download_file_ptr = NULL;
                            add_audit_log(ctx->db, ctx->user_id, "DOWNLOAD_ABORTED", ctx->current_download_server_path, ERR_CONNECTION_ERROR);
                        }
                        sm2_free_key_pair(&ctx->sm2_server_key);
                        free(ctx);
                    }
                }
            }
        }
    }

    free(events);
    // Cleanup (normally reached via signal handler)
    printf("Server main loop exited. Cleaning up...\n");
    if (g_epoll_fd != -1)
        close(g_epoll_fd);
    server_cleanup(g_server_fd); // Already sets g_server_fd to -1 if it closes it.
    destroy_server_thread_pool();
    close_database(g_db);

    return EXIT_SUCCESS;
}